import numpy as np
import random as rd

class HMM():
    def __init__(self, hidenStages, possiableObservasions, epoches =100):
        self.epoches = epoches
        self.hidenStages = hidenStages
        self.stagesIndex = {}
        self.observationIndex = {}
        self.possiableObservations = possiableObservasions
        self.stagesNum = len(hidenStages)
        self.observationNum = len(possiableObservasions)
        for i in range(self.stagesNum):
            self.stagesIndex[self.hidenStages[i]] = i
        for i in range(self.observationNum):
            self.observationIndex[self.possiableObservations[i]] = i
        
        self.convertMatrix = [[0.0 for i in range(self.stagesNum)] for j in range(self.stagesNum)]
        self.issueMatrix = [[0.0 for i in range(self.observationNum)] for j in range(self.stagesNum)]
        for i in range(self.stagesNum):
            remain = 1
            for j in range(self.stagesNum-1):
                rdm = rd.uniform(0, remain)
                remain = remain - rdm
                self.convertMatrix[i][j] = rdm
            self.convertMatrix[i][self.stagesNum-1] = remain
        for i in range(self.stagesNum):
            remain = 1
            for j in range(self.observationNum-1):
                rdm = rd.uniform(0, remain)
                remain = remain - rdm
                self.issueMatrix[i][j] = rdm
            self.issueMatrix[i][self.observationNum-1] = remain
    def EM(self, observationSeries):
        for s in range(self.epoches):
            length = len(observationSeries)
            stages = [None for i in range(length)]
            for i in range(length):
                maxp = 0.0
                for j in range(self.stagesNum):
                    if self.issueMatrix[j][self.observationIndex[observationSeries[i]]] >= maxp:
                        maxp = self.issueMatrix[j][self.observationIndex[observationSeries[i]]]
                        stages[i] = self.hidenStages[j]
            tmpConvertCounterMatrix = [[0.0 for i in range(self.stagesNum)] for j in range(self.stagesNum)]
            tmpIssueCounterMatrix = [[0.0 for i in range(self.observationNum)] for j in range(self.stagesNum)]
            for i in range(1, length):
                tmpConvertCounterMatrix[self.stagesIndex[stages[i-1]]][self.stagesIndex[stages[i]]] += 1
            for i in range(length):
                tmpIssueCounterMatrix[self.stagesIndex[stages[i]]][self.observationIndex[observationSeries[i]]] += 1
            for i in range(self.stagesNum):
                tmp = 0
                for j in range(self.stagesNum):
                    tmp += tmpConvertCounterMatrix[i][j]
                if tmp != 0:
                    for j in range(self.stagesNum):
                        tmpConvertCounterMatrix[i][j] = tmpConvertCounterMatrix[i][j] / tmp
                tmp = 0
                for j in range(self.observationNum):
                    tmp += tmpIssueCounterMatrix[i][j]
                if tmp != 0:
                    for j in range(self.observationNum):
                        tmpIssueCounterMatrix[i][j] = tmpIssueCounterMatrix[i][j] / tmp
            self.convertMatrix = tmpConvertCounterMatrix
            self.issueMatrix = tmpIssueCounterMatrix
    
    def setPossibilityMatrixs(self, convertModel, issueModel):
        self.convertMatrix = convertModel
        self.issueMatrix = issueModel
    
    def viterbi(self, observationSeries):
        length = len(observationSeries)
        tmpMatrix = [[0.0 for i in range(self.stagesNum)] for j in range(length)]
        traceBack = [[-1 for i in range(self.stagesNum)] for j in range(length)]
        for i in range(self.stagesNum):
            tmpMatrix[0][i] = self.issueMatrix[i][self.observationIndex[observationSeries[0]]]
        for i in range(1, length):
            for j in range(self.stagesNum):
                maxP = 0
                for k in range(self.stagesNum):
                    tmpP = tmpMatrix[i-1][k] * self.convertMatrix[k][j] * self.issueMatrix[j][self.observationIndex[observationSeries[i]]]
                    if tmpP > maxP:
                        maxP = tmpP
                        tmpMatrix[i][j] = tmpP
                        traceBack[i][j] = k
        maxPP = 0.0
        lastK = 0
        for i in range(self.stagesNum):
            if tmpMatrix[length-1][i] > maxPP:
                maxPP = tmpMatrix[length-1][i]
                lastK = i
        viterbiTrace = []
        for i in range(length):
            viterbiTrace.append(lastK)
            lastK = traceBack[length - 1 - i][lastK]
        viterbiTrace.reverse()
        self.viterbiTraceBackMatrix = traceBack
        self.viterbiProbabilityMatrix = tmpMatrix
        self.viterbiBestTrace = viterbiTrace

    def save(self):
        nparrConvert = np.array(self.convertMatrix)
        nparrIssue = np.array(self.issueMatrix)
        nparrViterbiTrace = np.array(self.viterbiBestTrace)
        np.savetxt("convertMatrix.txt", nparrConvert)
        np.savetxt("issueMatrix.txt", nparrIssue)
        np.savetxt("viterbiBestTrace.txt", nparrViterbiTrace)


hiddenStages = ["a", "b"]
possiableObservations = ["1", "2", "3"]
hmm = HMM(hiddenStages, possiableObservations, 100)
observationSeries = ["1", "2", "1", "1", "2", "2", "1", "1", "2", "2"]
convert = [[0.6, 0.4], [0.5, 0.5]]
issue = [[0.66667, 0.33333], [0.25, 0.75]]
hmm.setPossibilityMatrixs(convert, issue)
hmm.viterbi(observationSeries)
hmm.save()